import yaml

from core.connection import establish_connection
from core.query_set import QuerySet
from pattern.consume import c_torch_save_consumer
from pattern.traverse import q_node_contains_code, q_dfg_paths


# Load config
with open("config.yaml") as f:
    config = yaml.safe_load(f)
g = establish_connection(config["CONNECTION_STRING"])

# Query & Consume
q = QuerySet(g.V())

q.apply(q_node_contains_code, code="torch.save(", position="start")
q.apply(q_dfg_paths)

c_torch_save_consumer(q)
